from pydantic import BaseModel
from typing import List, Optional

class EmailData(BaseModel):
    text: str
    label: str  # "spam" or "ham"

class DatasetUpload(BaseModel):
    emails: List[EmailData]
    dataset_name: str

class DatasetInfo(BaseModel):
    dataset_name: str
    total_count: int
    spam_count: int
    ham_count: int
    
class TrainingConfig(BaseModel):
    dataset_name: str
    epochs: int = 3
    batch_size: int = 32
    learning_rate: float = 2e-5
    
class TrainingResponse(BaseModel):
    status: str
    message: str
    metrics: Optional[dict] = None 